package mongodb

import (
	"log"
	"sync"
	"time"

	"github.com/globalsign/mgo"

	"cmt1-admin/conf"
)

type MgoSession struct {
	mgoSession *mgo.Session
	db         string
	rw         *sync.RWMutex
}

func NewMgoSession(c *conf.MongoDB) (r *MgoSession, err error) {
	dialInfo := &mgo.DialInfo{
		Addrs:     []string{c.Addrs},
		Timeout:   time.Duration(c.Timeout) * time.Second,
		Source:    c.Source,
		Username:  c.Username,
		Password:  c.Password,
		PoolLimit: int(c.PoolLimit),
	}

	s, err := mgo.DialWithInfo(dialInfo)
	if err != nil {
		return
	}
	r = &MgoSession{mgoSession: s, db: c.DbName, rw: new(sync.RWMutex)}
	return
}

func (this *MgoSession) connect(collection string) (*mgo.Session, *mgo.Collection) {
	ms := this.mgoSession.Clone()
	c := ms.DB(this.db).C(collection)
	ms.SetMode(mgo.Monotonic, true)
	return ms, c
}

func (this *MgoSession) getDb(db string) (*mgo.Session, *mgo.Database) {
	ms := this.mgoSession.Clone()
	return ms, ms.DB(this.db)
}

func (this *MgoSession) IsEmpty(collection string) bool {
	ms, c := this.connect(collection)
	defer ms.Close()
	count, err := c.Count()
	if err != nil {
		log.Fatal(err)
	}
	return count == 0
}

func (this *MgoSession) Count(collection string, query interface{}) (int, error) {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Find(query).Count()
}

func (this *MgoSession) Insert(collection string, docs ...interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Insert(docs...)
}

func (this *MgoSession) FindOne(collection string, query, selector, result interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Find(query).Select(selector).One(result)
}

func (this *MgoSession) FindAll(collection string, query, selector, result interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Find(query).Select(selector).All(result)
}

func (this *MgoSession) FindAllSortLimit(collection string, query, selector, result interface{}, sorts []string, limit int) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Find(query).Select(selector).Sort(sorts...).Limit(limit).All(result)
}

func (this *MgoSession) FindPage(collection string, page, limit int, query, selector, result interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Find(query).Select(selector).Skip(page * limit).Limit(limit).All(result)
}

func (this *MgoSession) FindPageSort(collection string, page, limit int, query, selector interface{}, sorts []string, result interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Find(query).Select(selector).Sort(sorts...).Skip(page * limit).Limit(limit).All(result)
}

func (this *MgoSession) FindIter(collection string, query interface{}) *mgo.Iter {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Find(query).Iter()
}

func (this *MgoSession) Update(collection string, selector, update interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Update(selector, update)
}

func (this *MgoSession) Upsert(collection string, selector, update interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	_, err := c.Upsert(selector, update)
	return err
}

func (this *MgoSession) UpdateAll(collection string, selector, update interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	_, err := c.UpdateAll(selector, update)
	return err
}

func (this *MgoSession) Remove(collection string, selector interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	return c.Remove(selector)
}

func (this *MgoSession) RemoveAll(collection string, selector interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	_, err := c.RemoveAll(selector)
	return err
}

//insert one or multi documents
func (this *MgoSession) BulkInsert(collection string, docs ...interface{}) (*mgo.BulkResult, error) {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	bulk := c.Bulk()
	bulk.Insert(docs...)
	return bulk.Run()
}

func (this *MgoSession) BulkRemove(collection string, selector ...interface{}) (*mgo.BulkResult, error) {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	bulk := c.Bulk()
	bulk.Remove(selector...)
	return bulk.Run()
}

func (this *MgoSession) BulkRemoveAll(collection string, selector ...interface{}) (*mgo.BulkResult, error) {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	bulk := c.Bulk()
	bulk.RemoveAll(selector...)
	return bulk.Run()
}

func (this *MgoSession) BulkUpdate(collection string, pairs ...interface{}) (*mgo.BulkResult, error) {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	bulk := c.Bulk()
	bulk.Update(pairs...)
	return bulk.Run()
}

func (this *MgoSession) BulkUpdateAll(collection string, pairs ...interface{}) (*mgo.BulkResult, error) {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	bulk := c.Bulk()
	bulk.UpdateAll(pairs...)
	return bulk.Run()
}

func (this *MgoSession) BulkUpsert(collection string, pairs ...interface{}) (*mgo.BulkResult, error) {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	bulk := c.Bulk()
	bulk.Upsert(pairs...)
	return bulk.Run()
}

func (this *MgoSession) PipeAll(collection string, pipeline, result interface{}, allowDiskUse bool) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	var pipe *mgo.Pipe
	if allowDiskUse {
		pipe = c.Pipe(pipeline).AllowDiskUse()
	} else {
		pipe = c.Pipe(pipeline)
	}
	return pipe.All(result)
}

func (this *MgoSession) PipeOne(collection string, pipeline, result interface{}, allowDiskUse bool) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	var pipe *mgo.Pipe
	if allowDiskUse {
		pipe = c.Pipe(pipeline).AllowDiskUse()
	} else {
		pipe = c.Pipe(pipeline)
	}
	return pipe.One(result)
}

func (this *MgoSession) PipeIter(collection string, pipeline interface{}, allowDiskUse bool) *mgo.Iter {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	var pipe *mgo.Pipe
	if allowDiskUse {
		pipe = c.Pipe(pipeline).AllowDiskUse()
	} else {
		pipe = c.Pipe(pipeline)
	}

	return pipe.Iter()

}

func (this *MgoSession) Explain(collection string, pipeline, result interface{}) error {
	ms, c := this.connect(collection)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	pipe := c.Pipe(pipeline)
	return pipe.Explain(result)
}
func (this *MgoSession) GridFSCreate(prefix, name string) (*mgo.GridFile, error) {
	ms, d := this.getDb(this.db)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	gridFs := d.GridFS(prefix)
	return gridFs.Create(name)
}

func (this *MgoSession) GridFSFindOne(prefix string, query, result interface{}) error {
	ms, d := this.getDb(this.db)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	gridFs := d.GridFS(prefix)
	return gridFs.Find(query).One(result)
}

func (this *MgoSession) GridFSFindAll(prefix string, query, result interface{}) error {
	ms, d := this.getDb(this.db)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	gridFs := d.GridFS(prefix)
	return gridFs.Find(query).All(result)
}

func (this *MgoSession) GridFSOpen(prefix, name string) (*mgo.GridFile, error) {
	ms, d := this.getDb(this.db)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	gridFs := d.GridFS(prefix)
	return gridFs.Open(name)
}

func (this *MgoSession) GridFSRemove(prefix, name string) error {
	ms, d := this.getDb(this.db)
	defer ms.Close()
	this.rw.Lock()
	defer this.rw.Unlock()
	gridFs := d.GridFS(prefix)
	return gridFs.Remove(name)
}
